package cn.iocoder.yudao.framework.datasource.local;

import cn.hutool.core.util.StrUtil;

import java.util.Map;

/**
 * TODO
 *
 * @className: ThreadLocalContext
 * @author: tangkh
 * @date: 2023-09-19 17:19
 */
public class ThreadLocalContext {
    private static final ThreadLocal<ThreadLocalModel> threadLocal = new ThreadLocal();

    public ThreadLocalContext() {
    }

    public static ThreadLocalModel getLocalContext() {
        ThreadLocalModel threadLocalModel = (ThreadLocalModel)threadLocal.get();
        if (threadLocalModel == null) {
            threadLocalModel = new ThreadLocalModel();
            setLocalContext(threadLocalModel);
        }

        return (ThreadLocalModel)threadLocal.get();
    }

    public static void setLocalContext(ThreadLocalModel context) {
        threadLocal.set(context);
    }

    public static void removeLocalContext() {
        threadLocal.remove();
    }

    public static void removeAttachment() {
        removeAttachment((String)null);
    }

    public static void removeAttachment(String key) {
        ThreadLocalModel threadLocalModel = getLocalContext();
        Map<String, Object> attachments = threadLocalModel.getAttachments();
        if (StrUtil.isNotBlank(key)) {
            attachments.remove(key);
        } else {
            attachments.clear();
        }

    }

    public static Object getAttachment(String key) {
        ThreadLocalModel threadLocalModel = getLocalContext();
        Map<String, Object> attachments = threadLocalModel.getAttachments();
        return attachments.get(key);
    }

    public static void setAttachment(String key, Object value) {
        ThreadLocalModel threadLocalModel = getLocalContext();
        Map<String, Object> attachments = threadLocalModel.getAttachments();
        if (value == null) {
            attachments.remove(key);
        } else {
            attachments.put(key, value);
        }

    }
}
